from typing import Annotated, Optional
from pydantic import BaseModel, Field
import operator
from langgraph.graph import MessagesState
from langchain_core.messages import MessageLikeRepresentation
from typing_extensions import TypedDict

###################
# Structured Outputs
###################
class ConductResearch(BaseModel):
    """Call this tool to conduct research on a specific topic."""
    research_topic: str = Field(
        description="The topic to research. Should be a single topic, and should be described in high detail (at least a paragraph).",
    )

class ResearchComplete(BaseModel):
    """Call this tool to indicate that the research is complete."""

class Summary(BaseModel):
    summary: str
    key_excerpts: str

class ClarifyWithUser(BaseModel):
    need_clarification: bool = Field(
        description="Whether the user needs to be asked a clarifying question.",
    )
    question: str = Field(
        description="A question to ask the user to clarify the report scope",
    )
    verification: str = Field(
        description="Verify message that we will start research after the user has provided the necessary information.",
    )

class ResearchQuestion(BaseModel):
    research_brief: str = Field(
        description="A research question that will be used to guide the research.",
    )


###################
# State Definitions
###################

def override_reducer(current_value, new_value):
    """
    自定义状态合并函数，用于处理状态更新
    
    当新值是字典且包含type="override"时，直接使用新值覆盖当前值
    否则，使用operator.add将新值追加到当前值
    
    Args:
        current_value: 当前状态值
        new_value: 要合并的新值
        
    Returns:
        合并后的值
    """
    if isinstance(new_value, dict) and new_value.get("type") == "override":
        return new_value.get("value", new_value)
    else:
        return operator.add(current_value, new_value)
    
class AgentInputState(MessagesState):
    """InputState is only 'messages'"""

class AgentState(MessagesState):
    """
    深度研究代理的核心状态类
    
    继承自MessagesState，包含研究流程所需的各种状态字段
    用于在不同节点间传递和更新状态信息
    
    Attributes:
        supervisor_messages: 主管消息列表，使用override_reducer进行状态合并
        research_brief: 研究简报，可选字符串
        raw_notes: 原始笔记列表，使用override_reducer进行状态合并，初始为空列表
        notes: 处理后的笔记列表，使用override_reducer进行状态合并，初始为空列表
        final_report: 最终研究报告字符串
    """
    supervisor_messages: Annotated[list[MessageLikeRepresentation], override_reducer]
    research_brief: Optional[str]
    raw_notes: Annotated[list[str], override_reducer] = []
    notes: Annotated[list[str], override_reducer] = []
    final_report: str

class SupervisorState(TypedDict):
    """
    研究主管的状态类
    
    包含主管节点所需的状态信息
    
    Attributes:
        supervisor_messages: 主管消息列表，使用override_reducer进行状态合并
        research_brief: 研究简报字符串
        notes: 笔记列表，使用override_reducer进行状态合并，初始为空列表
        research_iterations: 研究迭代次数，初始为0
        raw_notes: 原始笔记列表，使用override_reducer进行状态合并，初始为空列表
    """
    supervisor_messages: Annotated[list[MessageLikeRepresentation], override_reducer]
    research_brief: str
    notes: Annotated[list[str], override_reducer] = []
    research_iterations: int = 0
    raw_notes: Annotated[list[str], override_reducer] = []

class ResearcherState(TypedDict):
    """
    研究员的状态类
    
    包含研究员节点所需的状态信息
    
    Attributes:
        researcher_messages: 研究员消息列表，使用operator.add进行状态合并
        tool_call_iterations: 工具调用迭代次数，初始为0
        research_topic: 研究主题字符串
        compressed_research: 压缩后的研究内容字符串
        raw_notes: 原始笔记列表，使用override_reducer进行状态合并，初始为空列表
    """
    researcher_messages: Annotated[list[MessageLikeRepresentation], operator.add]
    tool_call_iterations: int = 0
    research_topic: str
    compressed_research: str
    raw_notes: Annotated[list[str], override_reducer] = []

class ResearcherOutputState(BaseModel):
    """
    研究员输出状态类
    
    定义研究员节点的输出结构
    
    Attributes:
        compressed_research: 压缩后的研究内容字符串
        raw_notes: 原始笔记列表，使用override_reducer进行状态合并，初始为空列表
    """
    compressed_research: str
    raw_notes: Annotated[list[str], override_reducer] = []